package com.hyperf.spring.boot.autoconfigure;

import com.hyperf.core.shardingjdbc.PreciseShardingDatabaseAlgorithm;
import com.hyperf.core.shardingjdbc.PreciseShardingTableAlgorithm;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.standard.PreciseShardingAlgorithm;
import org.apache.shardingsphere.core.rule.ShardingRule;
import org.apache.shardingsphere.shardingjdbc.spring.datasource.SpringShardingDataSource;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.BeanFactory;
import org.springframework.beans.factory.BeanFactoryAware;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.sql.DataSource;

/**
 * shrdingjdbc分库分表自动装配
 *
 * @author lizhaoyang
 */
@SuppressWarnings("all")
@Slf4j
@Configuration
@ConditionalOnClass({DataSource.class, SpringShardingDataSource.class, PreciseShardingAlgorithm.class})
class ShardingAlgorithmAutoConfiguration implements BeanPostProcessor, BeanFactoryAware {

    private BeanFactory beanFactory;

    private ShardingRule shardingRule;

    @Bean("preciseShardingDatabaseAlgorithm")
    @ConditionalOnMissingBean
    public PreciseShardingDatabaseAlgorithm preciseShardingDatabaseAlgorithm() {
        return new PreciseShardingDatabaseAlgorithm();
    }

    @Bean("preciseShardingTableAlgorithm")
    @ConditionalOnMissingBean
    public PreciseShardingTableAlgorithm preciseShardingTableAlgorithm() {
        return new PreciseShardingTableAlgorithm();
    }

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (bean instanceof SpringShardingDataSource) {
            shardingRule = ((SpringShardingDataSource) bean).getRuntimeContext().getRule();
            PreciseShardingDatabaseAlgorithm shardingDatabaseAlgorithm = beanFactory.getBean(PreciseShardingDatabaseAlgorithm.class);
            shardingDatabaseAlgorithm.setShardingRule(shardingRule);
        }
        return bean;
    }

    @Override
    public void setBeanFactory(BeanFactory beanFactory) throws BeansException {
        this.beanFactory = beanFactory;
    }
}
